# 第一步：导入依赖的函数和库
# 第一步：导入依赖的函数和库
from elasticsearch_dsl import connections as es_connections
from pymilvus import connections as milvus_connections
from hybrid_search import hybrid_search  # 导入之前写的混合检索函数
from rag_test.pingan.pingan_text_embedding import text_embeddings  # 导入之前写的向量化函数
import os
from openai import OpenAI


# 第二步：先定义LLM调用函数（get_response）
def get_response(messages):
    """调用LLM生成回答（和之前讲的一致）"""
    client = OpenAI(
        # 注意：这里根据你的实际情况选择——本地模型或阿里云模型
        # 1. 本地模型（如Ollama部署的qwen，无需真实API密钥）
        # api_key="随便填（本地模型不校验）",
        base_url="http://localhost:11434/v1",

        # 2. 阿里云通义千问（需要替换成你的真实API密钥）
        api_key=os.getenv("DASHSCOPE_API_KEY") or "你的阿里云API密钥",
        #base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
    )

    completion = client.chat.completions.create(
        model="qwen:latest",  # 本地模型用这个；阿里云可用"qwen-plus"
        messages=messages,
        temperature=0.3  # 控制随机性，0.3更稳定
    )
    return completion.choices[0].message.content


# 第三步：定义完整的RAG流程函数（rag_pipeline）
def rag_pipeline(user_query):
    """
    完整RAG流程：用户提问→混合检索→构建上下文→LLM生成答案
    """
    # 1. 初始化ES和Milvus的连接（避免连接失败）
    # 初始化ES连接
    es_connections.create_connection(hosts=["http://localhost:9200"])
    milvus_connections.connect(host="localhost", port="19530")
    # 2. 执行混合检索：获取和问题最相关的3个文本片段（来自保险手册）
    search_results = hybrid_search(
        query=user_query,
        es_weight=0.3,  # 可调整的权重
        milvus_weight=0.7,
        top_k=3  # 只取最相关的3条，避免LLM信息过载
    )

    # 3. 把检索结果格式化成“上下文”（方便LLM参考，还能显示页码）
    context = ""
    for i, res in enumerate(search_results, 1):
        context += f"参考内容{i}（来自手册第{res['page_num']}页）：{res['text']}\n\n"

    # 4. 构建LLM的提示词（关键：限制LLM只能用参考内容回答，不编造）
    messages = [
        {
            "role": "system",
            "content": "你是平安保险手册专属助手，只能根据下面的参考内容回答用户问题。规则：1. 答案必须来自参考内容，不能编手册里没有的信息；2. 每个信息点要说明来自哪条参考内容（比如“根据参考内容1”）；3. 如果参考内容没相关信息，直接说“未找到平安保险手册中关于该问题的内容”。"
        },
        {
            "role": "user",
            "content": f"参考内容：\n{context}\n\n用户问题：{user_query}"
        }
    ]

    # 5. 调用LLM生成最终答案
    final_answer = get_response(messages)
    return final_answer


# 测试代码：直接运行这个文件时，能验证流程是否正常
if __name__ == "__main__":
    # 测试一个保险相关问题
    test_query = "平安福2024版的投保年龄是多少？"
    # 执行RAG流程
    answer = rag_pipeline(test_query)
    # 打印结果
    print(f"用户问题：{test_query}\n")
    print(f"系统回答：\n{answer}")